package com.itheima.consultant.controller;

import com.itheima.consultant.util.RedisUrlChecker;
import dev.langchain4j.community.store.embedding.redis.RedisEmbeddingStore;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.io.File;
import java.util.ArrayList;
import java.util.List;

import static java.util.Collections.singletonList;


/*
*
* 未测试
* 向量存储在redis的方式
* 存在embedding文件夹下，每个文件对应有一个或多个json片段存储向量
* json格式：
* {
* "file_name": "文件名，含后缀"
* "index": "分割出来时排第几"
* "vector": [向量数组，例如"0.687","0.646",…]
* "text": "分割的文字片段"
* "url": "/文件的绝对路径，注意前面多了个/"
*
* */

@RestController("/embedding")
public class EmbeddingController {
    @Autowired
    private EmbeddingModel embeddingModel;
    @Autowired
    private RedisEmbeddingStore redisEmbeddingStore;
    @Autowired
    private RedisTemplate redisTemplate;

    /*
    * （旧版本）这里是仅追加向量到redis，不会实现覆盖重写或去重
    * （新版本）这里是添加向量到redis，实现去重，不会覆盖重写，不过需要删除旧文件的向量才能重新添加
    *
    * 目前仅支持选择某文件或某文件夹，不支持多选（其实改用List<String>接收，下面遍历即可，这里就不改了）
    * */
    @GetMapping(value = "/add") //isFile()和isDirectory()
    public String addEmbedding(String DirectoryPath) {
        File path = new File(DirectoryPath);
        List<Document> documents = new ArrayList<>();

        String filepath;
        List<String> filename = new ArrayList<>();
        if (path.exists()) {
            Document document;
            if(path.isDirectory()) {
                filepath = path.getAbsolutePath();

                //documents=FileSystemDocumentLoader.loadDocuments(DirectoryPath); //方法1:要绝对路径（Document获取不到文件名和路径）
                //方法2:麻烦点，但能获取到文件名
                File[] files = path.listFiles(); //这里没用递归，有需要的小修一下即可
                if (files != null) {
                    for (File file : files) {
                        if (file.isFile()) {
                            filename.add(file.getName());
                            document = FileSystemDocumentLoader.loadDocument(path.toPath());
                            documents.add(document);
                        }
                    }
                }
            } else if(path.isFile()) {
                filepath = path.getParent();
                document = FileSystemDocumentLoader.loadDocument(path.toPath());
                documents.add(document);
            } else {
                return "unknow isDirectory or isFile";
            }
        } else {
            return "Directory does not exist";
        }

        //2.构建向量数据库操作对象  操作的是内存版本的向量数据库（因为InMemory数据都存在内存中不持久化，每次项目重启都要跑一遍向量化）
        //InMemoryEmbeddingStore store = new InMemoryEmbeddingStore();

        //构建文档分割器对象（依次分割段落、行、句子、词，直至最大容纳字符数）(每个片段最大容纳字符数，两个片段重叠的字符数-保证相邻片段的连贯性)
        DocumentSplitter ds = DocumentSplitters.recursive(500,100); //默认的文档分割器
        //3.构建一个EmbeddingStoreIngestor对象,完成文本数据切割,向量化, 存储
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                //.embeddingStore(store) //基于内存版本的非持久向量化
                .embeddingStore(redisEmbeddingStore) //基于redis缓存的持久向量化
                .documentSplitter(ds)
                .embeddingModel(embeddingModel) //向量模型
                .build();

        // 处理前检查文件是否已注入到向量中（只看是否同url且同文件名，不看同url同文件名时内容是否不同）（旧版本）
        RedisUrlChecker checker = new RedisUrlChecker(redisTemplate);
        for(int i=0;i<documents.size();i++) {
            String url="/"+filepath+"/"+filename.get(i); //好像是这样
            if (!checker.checkUrlExists(url)) { //不存在对应向量加入在redis的同目录同名文件
                ingestor.ingest(singletonList(documents.get(i)));
                System.out.println(url+"成功添加向量到redis");
            }
        }

//        ingestor.ingest(documents); //新版本
        return "success";
    }

    /*
    * 如果需要删除的，仿照上面的添加功能，改for(int i=0;i<documents.size();i++)里面的代码即可
    * */
}
